# -*- coding: utf-8 -*-                                                                                                             
import hashlib
import uuid
from string import join,split,strip
from sqlalchemy import create_engine, Table, Column, Integer, String, Boolean,Date, MetaData, select
from datetime import *


""" Class da base de dados em ORM"""

DATABASE_NAME = 'database/db.sqlite'

class Database:
    def __init__(self):
        self.connected = False
        

    def connect(self):
        self.engine = create_engine('sqlite:///' + DATABASE_NAME)
        self.connection = self.engine.connect()
        self.connected = True if self.connection else False
        self.metadata = MetaData()

        # Define the tables
        self.db_table_users = Table('users', self.metadata,
            Column('id', Integer, primary_key=True),
            Column('username', String, nullable=False),
            Column('password', String, nullable=False)
        )

        self.db_table_open_ports = Table('open_ports',self.metadata,
            Column('id',Integer,primary_key=True),
            Column('scanned_host',String,nullable=False),
            Column('open_port',String,nullable=False),
            Column('scan_date',Date, nullable=False),
            Column('username_id',String,nullable=False)
        )

        self.db_table_active_connections = Table('active_connections',self.metadata,
            Column('id',Integer,primary_key=True),
            Column('proto',String,nullable=False),
            Column('ip_src',String,nullable=False),
            Column('port_src',String,nullable=False),
            Column('ip_dest',String,nullable=False),
            Column('port_dest',String,nullable=False),
            Column('connection_status',String,nullable=False),
            Column('scan_date',Date,nullable=False),
            Column('username_id',String,nullable=False)
        )

        self.db_table_firewall_logs = Table('firewall_logs',self.metadata,
            Column('id',Integer,primary_key=True),
            Column('date',String,nullable=False),
            Column('type',String,nullable=False),
            Column('ip_src',String,nullable=False),
            Column('ip_dst',String,nullable=False),
            Column('proto',String,nullable=False),
            Column('username_id',String,nullable=False),
        )

        self.metadata.create_all(self.engine)

    def loginAccount(self,username,password):
        if not self.connected:
            return False
        s = select([self.db_table_users]).where(self.db_table_users.c.username == username)
        res = self.connection.execute(s)
        result = res.fetchall()
        res.close()
        if len(result) == 1:
            return self.decrypt_password(result[0]['password'],password)
        else:
            return False

    def createAccount(self,username,password):
        if not self.connected:
            return False

        s = select([self.db_table_users.c.username]).where(self.db_table_users.c.username == username)
        res = self.connection.execute(s)
        result = res.fetchall()
        res.close()
        if len(result) > 0:
            return False
        res = self.connection.execute(self.db_table_users.insert().values(username=username,password=self.crypt_password(password)))
        res.close()

        return True

    def addOpenPorts(self,host,port,username_id):
        if not self.connected:
            return False
        res = self.connection.execute(self.db_table_open_ports.insert().values(scanned_host=host,open_port=port,scan_date=datetime.now(),username_id=username_id))
        res.close()

        return True

    def getOpenPortsByHost(self,host,username_id):
        if not self.connected:
            return []

        s = select([self.db_table_open_ports]).where(self.db_table_open_ports.c.scanned_host == host).where(self.db_table_open_ports.c.username_id == username_id)
        res = self.connection.execute(s)
        result = res.fetchall()
        res.close()
        if len(result) > 0:
            return result
        else:
            return []

    def crypt_password(self,raw_password):
        salt = uuid.uuid4().hex
        return hashlib.sha256(salt.encode() + raw_password.encode()).hexdigest() + ':' + salt

    def decrypt_password(self,cypher_password, raw_password):
        password, salt = cypher_password.split(':')
        return password == hashlib.sha256(salt.encode() + raw_password.encode()).hexdigest()